import matplotlib.pyplot as plt
import matplotlib.patches as patches
from matplotlib.animation import FuncAnimation
import numpy as np

class SatelliteVisualizer:
    """卫星通信可视化器"""
    
    def __init__(self):
        self.fig = None
        self.ax1 = None  # 3D轨道图
        self.ax2 = None  # 地面轨迹图
        self.ax3 = None  # 链路预算图
        
        # 数据存储
        self.orbit_data = {'x': [], 'y': [], 'z': []}
        self.ground_track_data = {'lat': [], 'lon': []}
        self.link_data = {'time': [], 'snr': [], 'rx_power': []}
        
        # 地球参数
        self.earth_radius = 6371  # km
        
    def setup_plot(self, fig):
        """设置绘图"""
        self.fig = fig
        self.fig.clear()
        
        # 创建子图
        self.ax1 = self.fig.add_subplot(221, projection='3d')  # 3D轨道
        self.ax2 = self.fig.add_subplot(222)  # 地面轨迹
        self.ax3 = self.fig.add_subplot(212)  # 链路预算
        
        self._setup_3d_plot()
        self._setup_ground_track_plot()
        self._setup_link_budget_plot()
        
    def _setup_3d_plot(self):
        """设置3D轨道图"""
        self.ax1.clear()
        
        # 绘制地球
        u = np.linspace(0, 2 * np.pi, 50)
        v = np.linspace(0, np.pi, 50)
        x = self.earth_radius * np.outer(np.cos(u), np.sin(v))
        y = self.earth_radius * np.outer(np.sin(u), np.sin(v))
        z = self.earth_radius * np.outer(np.ones(np.size(u)), np.cos(v))
        
        self.ax1.plot_surface(x, y, z, alpha=0.3, color='blue')
        
        # 设置标签
        self.ax1.set_xlabel('X (km)')
        self.ax1.set_ylabel('Y (km)')
        self.ax1.set_zlabel('Z (km)')
        self.ax1.set_title('卫星轨道3D视图')
        
        # 设置坐标轴范围
        max_range = self.earth_radius * 1.5
        self.ax1.set_xlim([-max_range, max_range])
        self.ax1.set_ylim([-max_range, max_range])
        self.ax1.set_zlim([-max_range, max_range])
        
    def _setup_ground_track_plot(self):
        """设置地面轨迹图"""
        self.ax2.clear()
        
        # 绘制世界地图轮廓
        self.ax2.set_xlim([-180, 180])
        self.ax2.set_ylim([-90, 90])
        self.ax2.set_xlabel('经度 (°)')
        self.ax2.set_ylabel('纬度 (°)')
        self.ax2.set_title('卫星地面轨迹')
        self.ax2.grid(True)
        
        # 添加大陆轮廓（简化）
        self.ax2.axhline(y=0, color='gray', linestyle='--', alpha=0.5)
        self.ax2.axvline(x=0, color='gray', linestyle='--', alpha=0.5)
        
    def _setup_link_budget_plot(self):
        """设置链路预算图"""
        self.ax3.clear()
        
        self.ax3.set_xlabel('时间 (s)')
        self.ax3.set_ylabel('功率/信噪比 (dB)')
        self.ax3.set_title('链路预算分析')
        self.ax3.grid(True)
        self.ax3.legend()
        
    def update_plot(self, position, link_params):
        """更新绘图"""
        # 更新3D轨道图
        self._update_3d_orbit(position)
        
        # 更新地面轨迹
        self._update_ground_track(position)
        
        # 更新链路预算
        self._update_link_budget(link_params)
        
    def _update_3d_orbit(self, position):
        """更新3D轨道图"""
        # 清除之前的卫星位置
        if hasattr(self, 'satellite_point'):
            self.satellite_point.remove()
            
        # 绘制卫星位置
        self.satellite_point = self.ax1.scatter(
            [position['x']], [position['y']], [position['z']], 
            color='red', s=100, marker='o'
        )
        
        # 绘制卫星到地面站的连线
        if hasattr(self, 'link_line'):
            self.link_line.remove()
            
        # 地面站位置（假设在赤道）
        gs_x = self.earth_radius
        gs_y = 0
        gs_z = 0
        
        self.link_line = self.ax1.plot(
            [position['x'], gs_x], 
            [position['y'], gs_y], 
            [position['z'], gs_z], 
            'g--', alpha=0.7
        )[0]
        
    def _update_ground_track(self, position):
        """更新地面轨迹图"""
        # 添加新的地面轨迹点
        self.ground_track_data['lat'].append(position['latitude'])
        self.ground_track_data['lon'].append(position['longitude'])
        
        # 限制数据点数量
        max_points = 100
        if len(self.ground_track_data['lat']) > max_points:
            self.ground_track_data['lat'] = self.ground_track_data['lat'][-max_points:]
            self.ground_track_data['lon'] = self.ground_track_data['lon'][-max_points:]
        
        # 清除之前的轨迹
        self.ax2.clear()
        self._setup_ground_track_plot()
        
        # 绘制轨迹
        if len(self.ground_track_data['lat']) > 1:
            self.ax2.plot(
                self.ground_track_data['lon'], 
                self.ground_track_data['lat'], 
                'r-', linewidth=2, label='地面轨迹'
            )
            
        # 绘制当前位置
        self.ax2.scatter(
            [position['longitude']], 
            [position['latitude']], 
            color='red', s=100, marker='o', label='当前位置'
        )
        
        # 绘制覆盖区域
        coverage_radius = 10  # 简化的覆盖半径
        circle = patches.Circle(
            (position['longitude'], position['latitude']), 
            coverage_radius, 
            fill=False, 
            color='green', 
            linestyle='--',
            label='覆盖区域'
        )
        self.ax2.add_patch(circle)
        
        self.ax2.legend()
        
    def _update_link_budget(self, link_params):
        """更新链路预算图"""
        # 添加新的数据点
        current_time = len(self.link_data['time']) * 0.1  # 假设每0.1秒一个点
        self.link_data['time'].append(current_time)
        self.link_data['snr'].append(link_params['snr'])
        self.link_data['rx_power'].append(link_params['rx_power'])
        
        # 限制数据点数量
        max_points = 50
        for key in self.link_data:
            if len(self.link_data[key]) > max_points:
                self.link_data[key] = self.link_data[key][-max_points:]
        
        # 清除之前的图
        self.ax3.clear()
        
        # 绘制SNR
        self.ax3.plot(
            self.link_data['time'], 
            self.link_data['snr'], 
            'b-', linewidth=2, label='信噪比 (SNR)'
        )
        
        # 绘制接收功率
        self.ax3.plot(
            self.link_data['time'], 
            self.link_data['rx_power'], 
            'r-', linewidth=2, label='接收功率 (dBm)'
        )
        
        # 添加阈值线
        self.ax3.axhline(y=10, color='g', linestyle='--', label='最小SNR要求')
        
        self.ax3.set_xlabel('时间 (s)')
        self.ax3.set_ylabel('功率/信噪比 (dB)')
        self.ax3.set_title('链路预算分析')
        self.ax3.grid(True)
        self.ax3.legend()
        
    def reset(self):
        """重置可视化数据"""
        self.orbit_data = {'x': [], 'y': [], 'z': []}
        self.ground_track_data = {'lat': [], 'lon': []}
        self.link_data = {'time': [], 'snr': [], 'rx_power': []}
        
        if self.fig is not None:
            self.setup_plot(self.fig)